import gradio as gr
import time

from UserAcess import auth

from langchain_core.output_parsers import StrOutputParser
from ChatbotWithRetrieval import ChatbotWithRetrieval
from logger import setup_logger

base_dir = 'C:\\Users\Yang.Shen9\PycharmProjects\goudan-llm\OneFlower'
llm = ChatbotWithRetrieval(base_dir)

# 获取日志记录器对象
logger = setup_logger()


def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)


messageType = 0

output_parser = StrOutputParser()


def add_message(history, message):
    for x in message["files"]:
        history.append(((x,), None))
        store = llm.store
        store.load_file(x)
    if message["text"] != '':
        history.append((message["text"], None))
    return history, gr.MultimodalTextbox(value=None, interactive=False)


def bot(history):
    question = history[-1][0]
    if type(question) == tuple:
        logger.info("question:文件上传")
        yield from retMsg(history, "上传成功")
        logger.info("response:上传成功")
        return
    if question is None:
        return
    logger.info("question:" + question)
    chat_chain = llm.qa
    historyInfo = get_history(history)
    logger.info("gptHistoryInfo:" + historyInfo)
    response = chat_chain.invoke({"history": historyInfo, "input": question})["answer"]

    logger.info("response:" + response)
    yield from retMsg(history, response)


def_message = "你好，我是易速鲜的客服狗蛋儿，我可以为你提供鲜花及其公司员工的信息。"


def get_history(history):
    history = history[-2]
    question = history[0]
    if question is None:
        return ''
    answer = history[1]
    if answer == '上传成功' or answer == def_message:
        return ''
    history_question = """history question:""" + question
    history_answer = """history answer:""" + answer
    history = history_question + history_answer
    return history


def retMsg(history, response):
    history[-1][1] = ''
    for character in response:
        history[-1][1] += character
        time.sleep(0.05)
        yield history


message = [(None, def_message)]
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(
        message,
        elem_id="chatbot",
        bubble_full_width=False
    )

    chat_input = gr.MultimodalTextbox(interactive=True, file_types=["pdf", "txt", "docx", "doc"],
                                      placeholder="Enter message or upload file...", show_label=False)

    chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
    bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
    bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

    chatbot.like(print_like_dislike, None, None)

demo.queue()

if __name__ == "__main__":
    # demo.launch(auth=("admin", "pass1234"))
    demo.launch(auth=auth)
